! ==============================================================================
! This file is part of CUPP (Community Unified Physics Package) since 2023.
! ==============================================================================

module tracer_mod

  use cupp_const_mod
  use tracer_types_mod

  implicit none

  private

  public tracer_init
  public tracer_final
  public tracer_add
  public tracer_allocate
  public tracer_get_idx
  public tracer_get_array
  public tracer_get_array_qm
  public tracer_calc_qm
  public ntracers
  public ntracers_water
  public idx_qv
  public idx_qc, idx_nc
  public idx_qi, idx_ni
  public idx_qr, idx_nr
  public idx_qs, idx_ns
  public idx_qg
  public idx_qh
  public idx_qo3
  public idx_qso2
  public tracer_names
  public tracer_long_names
  public tracer_units
  public tracer_types
  public tracers

  interface tracer_get_array
    module procedure tracer_get_array_idx
    module procedure tracer_get_array_name
    module procedure tracer_get_array_all
  end interface tracer_get_array

contains

  subroutine tracer_init()

    call tracer_final()

    allocate(tracer_names     (100)); tracer_names      = 'N/A'
    allocate(tracer_long_names(100)); tracer_long_names = 'N/A'
    allocate(tracer_units     (100)); tracer_units      = 'kg kg-1'
    allocate(tracer_types     (100)); tracer_types      = 0

  end subroutine tracer_init

  subroutine tracer_final()

    ntracers       = 0
    ntracers_water = 0
    idx_qv         = 0
    idx_qc         = 0
    idx_nc         = 0
    idx_qi         = 0
    idx_ni         = 0
    idx_qr         = 0
    idx_nr         = 0
    idx_qs         = 0
    idx_ns         = 0
    idx_qg         = 0
    idx_qh         = 0
    idx_qo3        = 0

    if (allocated(tracer_names     )) deallocate(tracer_names     )
    if (allocated(tracer_long_names)) deallocate(tracer_long_names)
    if (allocated(tracer_units     )) deallocate(tracer_units     )
    if (allocated(tracer_types     )) deallocate(tracer_types     )
    if (allocated(tracers          )) deallocate(tracers          )

  end subroutine tracer_final

  subroutine tracer_add(name, long_name, units, type)

    character(*), intent(in) :: name
    character(*), intent(in) :: long_name
    character(*), intent(in), optional :: units
    integer, intent(in), optional :: type

    integer i
    logical found

    ntracers = ntracers + 1
    tracer_names(ntracers)      = name
    tracer_long_names(ntracers) = long_name
    if (present(units)) tracer_units(ntracers) = units
    if (present(type )) tracer_types(ntracers) = type

  end subroutine tracer_add

  subroutine tracer_allocate(mesh)

    type(cupp_mesh_type), intent(in) :: mesh(:)

    integer iblk, m

    ! Set tracer indices.
    do m = 1, ntracers
      select case (tracer_names(m))
      case ('qv', 'Q')
        idx_qv    = m; ntracers_water = ntracers_water + 1
      case ('qc', 'CLDLIQ')
        idx_qc    = m; ntracers_water = ntracers_water + 1
      case ('nc', 'NUMLIQ')
        idx_nc    = m
      case ('qi', 'CLDICE')
        idx_qi    = m; ntracers_water = ntracers_water + 1
      case ('ni', 'NUMICE')
        idx_ni    = m
      case ('qr', 'RAINQM')
        idx_qr    = m; ntracers_water = ntracers_water + 1
      case ('qs', 'SNOWQM')
        idx_qs    = m; ntracers_water = ntracers_water + 1
      case ('qg')
        idx_qg    = m; ntracers_water = ntracers_water + 1
      case ('qh')
        idx_qh    = m; ntracers_water = ntracers_water + 1
      case ('qo3')
        idx_qo3   = m
      case ('qso2', 'SO2')
        idx_qso2  = m
      end select
    end do

    ! Allocate tracer arrays for each block.
    allocate(tracers(size(mesh)))
    do iblk = 1, size(mesh)
      call tracers(iblk)%init(mesh(iblk), ntracers)
    end do

  end subroutine tracer_allocate

  pure integer function tracer_get_idx(name) result(res)

    character(*), intent(in) :: name

    integer m

    do m = 1, ntracers
      if (name == tracer_names(m)) then
        res = m
        return
      end if
    end do
    res = 0

  end function tracer_get_idx

  subroutine tracer_get_array_idx(iblk, idx, q, file, line)

    integer, intent(in) :: iblk
    integer, intent(in) :: idx
    real(r8), intent(out), pointer :: q(:,:)
    character(*), intent(in) :: file
    integer, intent(in) :: line

    if (idx < 1) then
      q => null()
    else
      q => tracers(iblk)%q(:,:,idx)
    end if

  end subroutine tracer_get_array_idx

  subroutine tracer_get_array_name(iblk, name, q, file, line)

    integer, intent(in) :: iblk
    character(*), intent(in) :: name
    real(r8), intent(out), pointer :: q(:,:)
    character(*), intent(in) :: file
    integer, intent(in) :: line

    integer idx

    idx = tracer_get_idx(name)
    call tracer_get_array(iblk, idx, q, file, line)

  end subroutine tracer_get_array_name

  subroutine tracer_get_array_all(iblk, q)

    integer, intent(in) :: iblk
    real(r8), intent(out), pointer :: q(:,:,:)

    if (.not. allocated(tracers)) return

    q => tracers(iblk)%q

  end subroutine tracer_get_array_all

  subroutine tracer_calc_qm(iblk)

    integer, intent(in) :: iblk

    integer icol, ilev

    if (.not. allocated(tracers)) return
    if (.not. allocated(tracers(iblk)%qm)) return

    associate (mesh    => tracers(iblk)%mesh, &
               tracers => tracers(iblk)     , &
               q       => tracers(iblk)%q   , & ! in
               qm      => tracers(iblk)%qm  )   ! out
    if (idx_qv > 0) then
      do ilev = 1, mesh%nlev
        do icol = 1, mesh%ncol
          qm(icol,ilev) = q(icol,ilev,idx_qv)
        end do
      end do
    end if
    if (idx_qc > 0) then
      do ilev = 1, mesh%nlev
        do icol = 1, mesh%ncol
          qm(icol,ilev) = qm(icol,ilev) + q(icol,ilev,idx_qc)
        end do
      end do
    end if
    if (idx_qi > 0) then
      do ilev = 1, mesh%nlev
        do icol = 1, mesh%ncol
          qm(icol,ilev) = qm(icol,ilev) + q(icol,ilev,idx_qi)
        end do
      end do
    end if
    if (idx_qr > 0) then
      do ilev = 1, mesh%nlev
        do icol = 1, mesh%ncol
          qm(icol,ilev) = qm(icol,ilev) + q(icol,ilev,idx_qr)
        end do
      end do
    end if
    if (idx_qs > 0) then
      do ilev = 1, mesh%nlev
        do icol = 1, mesh%ncol
          qm(icol,ilev) = qm(icol,ilev) + q(icol,ilev,idx_qs)
        end do
      end do
    end if
    if (idx_qg > 0) then
      do ilev = 1, mesh%nlev
        do icol = 1, mesh%ncol
          qm(icol,ilev) = qm(icol,ilev) + q(icol,ilev,idx_qg)
        end do
      end do
    end if
    if (idx_qh > 0) then
      do ilev = 1, mesh%nlev
        do icol = 1, mesh%ncol
          qm(icol,ilev) = qm(icol,ilev) + q(icol,ilev,idx_qh)
        end do
      end do
    end if
    end associate

  end subroutine tracer_calc_qm

  subroutine tracer_get_array_qm(iblk, qm)

    integer, intent(in) :: iblk
    real(r8), intent(out), pointer :: qm(:,:)

    if (.not. allocated(tracers)) return
    if (.not. allocated(tracers(iblk)%qm)) return

    qm => tracers(iblk)%qm

  end subroutine tracer_get_array_qm

end module tracer_mod
